"""
python fujian1_ArimaAfterSmooth.py


数据预处理:

    考虑对 inventory 列进行差分以使数据平稳（如果尚未平稳）。您可以使用 df['inventory'].diff().dropna() 来差分数据。
    如果数据存在异常值，考虑去除或替换它们。

使用更简单的模型:

    尝试使用较低阶数的 ARIMA 模型，比如 ARIMA(1,0,1) 或 ARIMA(2,0,0)，并观察是否能成功拟合。
    也可以尝试其他时间序列模型，比如季节性分解或指数平滑等。
"""
import os
import json
import pandas as pd
from pmdarima import auto_arima

# 设置输入和输出目录（使用相对路径）
input_dir = r"./fujian/fujian1/spline_then_smooth"
output_dir = r"./fujian/fujian1/smooth_then_arima"

# 创建输出目录（如果不存在）
os.makedirs(output_dir, exist_ok=True)

# 遍历目录中的所有 JSON 文件
for filename in os.listdir(input_dir):
    if filename.endswith('.json'):
        file_path = os.path.join(input_dir, filename)

        # 读取 JSON 文件
        with open(file_path, 'r') as f:
            data = json.load(f)

        # 将数据转换为 DataFrame
        df = pd.DataFrame(data)
        df['date'] = pd.to_datetime(df['date'])  # 转换日期格式
        df.set_index('date', inplace=True)

        # 确保有 12 个数据点进行训练
        if len(df) >= 12:
            # 填充缺失值（如果有）
            df['inventory'] = df['inventory'].ffill()  # 向前填充缺失值


            # 尝试 ARIMA(1,0,1) 或 ARIMA(2,0,0)
            try:
                model = auto_arima(df['inventory'], seasonal=False, stepwise=True, trace=True, suppress_warnings=True)
            except Exception as e:
                print(f"Failed to fit ARIMA model for {filename}: {e}")
                continue

            # 进行预测，预测未来 3 个数据点
            forecast, conf_int = model.predict(n_periods=3, return_conf_int=True)

            # 生成未来日期
            forecast_dates = pd.date_range(start=df.index[-1] + pd.offsets.MonthEnd(), periods=3, freq='ME')

            # 创建预测结果的 DataFrame
            forecast_result = pd.DataFrame({'date': forecast_dates, 'inventory': forecast})

            # 合并原始数据和预测数据
            result_df = pd.concat([df.reset_index(), forecast_result], ignore_index=True)

            # 转换日期为字符串格式以便 JSON 序列化
            result_df['date'] = result_df['date'].dt.strftime('%Y-%m-%dT%H:%M:%S.%f')

            # 转换回 JSON 格式并输出
            output_path = os.path.join(output_dir, filename)
            result_json = result_df.to_dict(orient='records')

            with open(output_path, 'w') as f:
                json.dump(result_json, f, indent=4)

            print(f"Processed {filename} and saved results to {output_path}")
        else:
            print(f"Not enough data points in {filename} for ARIMA model.")
